from PIL import Image, ImageSequence, ImageOps
import torch
import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
from io import BytesIO
import os
import numpy as np
import urllib3

# 禁用 SSL 验证警告
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
# 默认请求头，模拟浏览器访问，防止部分服务器拒绝
DEFAULT_HEADERS = {
    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64)"
}

def pil2tensor(img):
    output_images = []
    output_masks = []
    for i in ImageSequence.Iterator(img):
        i = ImageOps.exif_transpose(i)
        if i.mode == 'I':
            i = i.point(lambda i: i * (1 / 255))
        image = i.convert("RGB")
        image = np.array(image).astype(np.float32) / 255.0
        image = torch.from_numpy(image)[None,]
        if 'A' in i.getbands():
            mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0
            mask = 1. - torch.from_numpy(mask)
        elif i.mode == 'P' and 'transparency' in i.info:
            mask = np.array(i.convert('RGBA').getchannel('A')).astype(np.float32) / 255.0
            mask = 1. - torch.from_numpy(mask)
        else:
            mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu")
        output_images.append(image)
        output_masks.append(mask.unsqueeze(0))

    if len(output_images) > 1:
        output_image = torch.cat(output_images, dim=0)
        output_mask = torch.cat(output_masks, dim=0)
    else:
        output_image = output_images[0]
        output_mask = output_masks[0]

    return (output_image, output_mask)


def load_image(url: str, retries: int = 3, timeout: float = 10.0):
    """
    支持本地路径或 URL 加载图片。
    - 本地路径优先
    - 网络加载支持域名/IP, SSL 验证跳过, 重试机制
    - 忽略环境代理，确保直接访问
    - 添加浏览器 User-Agent 头，防止 403 拒绝
    - 放宽 Content-Type 检查，允许 application/octet-stream
    返回 PIL Image 对象及文件名
    """
    source = url.strip()
    filename = os.path.basename(source.split('?', 1)[0]) or 'downloaded_image'
    # 本地文件
    if os.path.exists(source):
        img = Image.open(source)
    else:
        # 网络加载
        session = requests.Session()
        session.trust_env = False
        session.headers.update(DEFAULT_HEADERS)
        retry_strategy = Retry(
            total=retries,
            backoff_factor=0.3,
            status_forcelist=[429, 500, 502, 503, 504],
            allowed_methods=["GET"]
        )
        adapter = HTTPAdapter(max_retries=retry_strategy)
        session.mount("https://", adapter)
        session.mount("http://", adapter)

        resp = session.get(source, timeout=timeout, verify=False)
        resp.raise_for_status()

        # 放宽 Content-Type 检查
        ctype = resp.headers.get('Content-Type', '').lower()
        # 部分服务器将图片标记为 application/octet-stream，也允许
        if ctype and not (ctype.startswith('image/') or 'octet-stream' in ctype):
            raise RuntimeError(f"URL does not point to an image, content-type: {ctype}")

        data = resp.content
        try:
            img = Image.open(BytesIO(data))
        except Exception:
            # 尝试强制以 RGB 打开
            img = Image.open(BytesIO(data)).convert("RGB")

    return img, filename


class JkhaLoadImage:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {"url": ("STRING", {"multiline": False})},
            "optional": {
                "retries": ("INT", {"default": 3, "min": 0, "max": 10}),
                "timeout": ("FLOAT", {"default": 10.0, "min": 1.0, "max": 60.0})
            }
        }

    RETURN_TYPES = ("IMAGE", "MASK")
    FUNCTION = "load"
    CATEGORY = "image"
    DESCRIPTION = (
        "京科华安加载图片节点V4,从url加载图片"
    )

    def load(self, url, retries=3, timeout=10.0):
        img, name = load_image(url, retries=retries, timeout=timeout)
        img_tensor, mask_tensor = pil2tensor(img)
        return img_tensor, mask_tensor


if __name__ == "__main__":
    img, name = load_image("")
    img_out, mask_out = pil2tensor(img)
